//------------------------------------------------------------------------------
// <copyright file="WinsockStringSocket.cpp" company="Microsoft">
//     Copyright (c) Microsoft Corporation.  All rights reserved.
// </copyright>
//------------------------------------------------------------------------------
#include "stdafx.h"

#ifndef _WINRT_DLL

#include "WinsockStringSocket.h"
#include <ws2ipdef.h>
#include <assert.h>

using namespace Microsoft::Internal::GamesTest::Rpc::Server;
using namespace concurrency;

#define HRESULT_FROM_WSA HRESULT_FROM_WIN32(WSAGetLastError())

WinsockStringSocket::WinsockStringSocket(SOCKET socket) :
    m_socket(socket), 
    m_dataReceivedCallback(nullptr), 
    m_clientDisconnectedCallback(nullptr),
    m_receiving(false),
    m_workerThread(nullptr),
    m_dataReceivedTaskInitialized(false)
{
}

WinsockStringSocket::~WinsockStringSocket(void)
{
    Shutdown();
}

HRESULT WinsockStringSocket::StartListening(DataReceivedCallback dataReceivedCallback, SocketDisconnectedCallback clientDisconnectedCallback)
{
    HRESULT hr = S_OK;
    m_dataReceivedCallback = dataReceivedCallback;
    m_clientDisconnectedCallback = clientDisconnectedCallback;

    if(!m_receiving)
    {
        m_receiving = true;

        DWORD threadId;
        m_workerThread = CreateThread(
            nullptr,
            0,			// default stack size
            (LPTHREAD_START_ROUTINE)ThreadStart,
            this,
            0,			// start immediately
            &threadId);

        if(m_workerThread == NULL)
        {
            RpcPrintMessage(L"WinsockStringSocket create thread failed");
            hr = HRESULT_FROM_WIN32(GetLastError());
            m_receiving = false;
            goto end;
        }
    }

end:
    return hr;
}

void WinsockStringSocket::Shutdown()
{
    m_dataReceivedCallback = nullptr;
    m_clientDisconnectedCallback = nullptr;

    // Closing the socket forces blocking calls like send/recv to unblock
    CloseSocket();

    m_receiving = false;

    if(m_workerThread != INVALID_HANDLE_VALUE)
    {
        WaitForSingleObject(m_workerThread, INFINITE);
        CloseHandle(m_workerThread);
        m_workerThread = INVALID_HANDLE_VALUE;
    }
}

DWORD WinsockStringSocket::Process()
{
    while(m_receiving)
    {
        int size;
        if(!Read7BitEncodedInt(size))
        {
            RaiseSocketDisconnectedEvent(m_clientDisconnectedCallback, this);
            break;
        }

        if(!ReadString(size))
        {
            RaiseSocketDisconnectedEvent(m_clientDisconnectedCallback, this);
            break;
        }		
    }

    return 0;
}

std::wstring WinsockStringSocket::GetIpAddress()
{
    SOCKADDR remoteSockAddr;
    int length = (int)sizeof(remoteSockAddr);
    int result = getpeername(m_socket, &remoteSockAddr, &length);

    if(result == SOCKET_ERROR)
    {
        RpcPrintMessage(L"Failed to retrieve socket name.");
        RpcPrintError(HRESULT_FROM_WSA);
        return nullptr;
    }

    TCHAR buffer[22]; // "xxx.xxx.xxx.xxx:xxxxx\0"
    DWORD len = _countof(buffer);

    result = WSAAddressToString(&remoteSockAddr, sizeof(remoteSockAddr), nullptr, buffer, &len);

    if(result == SOCKET_ERROR)
    {
        RpcPrintMessage(L"Failed to retrieve address as a string.");
        RpcPrintError(HRESULT_FROM_WSA);
        return nullptr;
    }

    return std::wstring(buffer);
}

// Turn compile optimization off for this function in order to use the Mock functionality in tests
// The reason we have created these RecvImpl / SendImpl functions is so we can mock recv in a different 
// Way when called to receive the length vs. receive the content
// Mock10 framework has limitations on filtered Mocking (Mock x only when called by y) when code is optimized so we had to take this approach
#pragma optimize( "", off)
int WinsockStringSocket::ReadByteRecvImpl(char* buffer, int bufferLength)
{
    return ::recv(m_socket, buffer, bufferLength, 0);
}

bool WinsockStringSocket::ReadByte(BYTE & byte)
{
    const int recvBufLength = 1;
    char recvbuf[recvBufLength];
    recvbuf[0] = 0;

    int result = ReadByteRecvImpl(recvbuf, recvBufLength);

    if (result > 0) 
    {
#ifdef _DEBUG
        const int bufferLength = 50;
        wchar_t buffer[bufferLength];		
        swprintf_s(buffer, bufferLength, L"Bytes received: %d", (int)recvbuf[0]);
        RpcPrintMessage(buffer);
#endif
    }
    else if(result == 0)
    {
        RpcPrintMessage(L"Failed to recv byte, socket closed");
        RpcPrintError(HRESULT_FROM_WSA);
        return false;
    }
    else if (result == SOCKET_ERROR)
    {		
        RpcPrintMessage(L"Receive byte failed, error");
        RpcPrintError(HRESULT_FROM_WSA);
        return false;
    }

    byte = recvbuf[0];

    return true;
}

// This function came pretty much straight from the .NET framework's BinaryReader class.
// https://codesearch/FileView.aspx?filename=%5c%5ctk5codeindex%5cdotnet%5c8.0%5cdevdiv%5cdepot%5cdevdiv%5creleases%5cwhidbey%5credbits%5cndp%5cclr%5csrc%5cbcl%5csystem%5cio%5cbinaryreader.cs%5c1%5cbinaryreader.cs&term="binaryreader"
bool WinsockStringSocket::Read7BitEncodedInt(int & i) 
{
    // Read out an Int32 7 bits at a time.  The high bit 
    // of the byte when on means to continue reading more bytes.
    int count = 0;
    int shift = 0;
    BYTE b = 0;

    do 
    {
        // Check for a corrupted stream.  Read a max of 5 bytes. 
        // In a future version, add a DataFormatException. 
        if (shift == 5 * 7)  // 5 bytes max per Int32, shift += 7
            return false;

        if(!ReadByte(b))
        {
            return false;
        }
        count |= (b & 0x7F) << shift; 
        shift += 7;
    } while ((b & 0x80) != 0); 

    i = count; 

    return true;
}

// This function came straight from the .NET framework's BinaryWriter class.
// https://codesearch/FileView.aspx?filename=%5c%5ctk5codeindex%5cdotnet%5c8.0%5cdevdiv%5cdepot%5cdevdiv%5creleases%5cwhidbey%5credbits%5cndp%5cclr%5csrc%5cbcl%5csystem%5cio%5cbinarywriter.cs%5c1%5cbinarywriter.cs&term="binarywriter"
int WinsockStringSocket::Write7BitEncodedInt(char buffer[5], int value) 
{
    // Write out an int 7 bits at a time.  The high bit of the byte, 
    // when on, tells reader to continue reading more bytes. 
    UINT32 v = (UINT32) value;   // support negative numbers
    int offset = 0;
    while (v >= 0x80) 
    { 
        assert(offset < 4);
        buffer[offset] = (BYTE)(v | 0x80);
        v >>= 7;
        offset++;
    }
    buffer[offset] = (BYTE)v;

    return ++offset;
}

// Turn compile optimization off for this function in order to use the Mock functionality in tests
// The reason we have created these RecvImpl / SendImpl functions is so we can mock send in a different 
// Way when called to receive the length vs. receive the content
// Mock10 framework has limitations on filtered Mocking (Mock x only when called by y) when code is optimized so we had to take this approach
// The function size must be at least 6 bytes in order to mock this function (as specified in Mock10 documentation)
#pragma optimize( "", off)
int WinsockStringSocket::DataLengthInBytesSendImpl(char* buffer, int length)
{
    int result = ::send(m_socket, buffer, length, 0);
    return result;
}

HRESULT WinsockStringSocket::SendDataLengthInBytes(int dataLength)
{
    HRESULT hr = S_OK;
    char lengthEncodedInt[5];

    // Send length of the multibyte buffer
    int intLength = Write7BitEncodedInt(lengthEncodedInt, dataLength);
    int result = 0;
    int bytesSent = 0;
    do
    {
        result = DataLengthInBytesSendImpl(lengthEncodedInt + bytesSent, intLength - bytesSent);
        if(result == SOCKET_ERROR)
        {
            RpcPrintMessage(L"socket send failed");
            hr = HRESULT_FROM_WSA;
            goto end;
        }

        bytesSent += result;
    }
    while(bytesSent < intLength);
end:
    return hr;	
}

// Turn compile optimization off for this function in order to use the Mock functionality in tests
// The reason we have created these RecvImpl / SendImpl functions is so we can mock send in a different 
// Way when called to receive the length vs. receive the content
// Mock10 framework has limitations on filtered Mocking (Mock x only when called by y) when code is optimized so we had to take this approach
// Turn compile optimization off for this function in order to use the Mock functionality in tests
#pragma optimize( "", off)
int WinsockStringSocket::DataContentSendImpl(char* buffer, int length)
{
    return ::send(m_socket, buffer, length, 0);
}

HRESULT WinsockStringSocket::SendDataContent(char* buffer, int bufferLength)
{
    HRESULT hr = S_OK;

    // Send multibyte buffer 
    int result = 0;
    int bytesSent = 0;

    do
    {
        result = DataContentSendImpl(buffer + bytesSent, bufferLength - bytesSent);
        if(result == SOCKET_ERROR)
        {
            RpcPrintMessage(L"Write string SOCKET_ERROR");
            hr = HRESULT_FROM_WSA;
            goto end;
        }

        bytesSent += result;
    }
    while(bytesSent < bufferLength);
end:
    return hr;
}

GAMESTEST_ANNOTATE_HRESULT_METHOD HRESULT WinsockStringSocket::SendData(const std::wstring& data)
{
    concurrency::critical_section::scoped_lock lock(s_synchronizationContext);

    HRESULT hr = S_OK;

    // Initialize buffers
    static const int BUFFER_SIZE = 512;
    char buffer[BUFFER_SIZE];
    char* pBuffer = buffer;

    const wchar_t* strData = data.c_str();	

	// Get length of buffer required to convert wide to UTF8
    int multiByteLength = WideCharToMultiByte(CP_UTF8, 0, strData, (int)data.length() + 1, NULL, 0, NULL, NULL);

    if(multiByteLength > BUFFER_SIZE)
    {
		// Allocate buffer to hold UTF8 result of the conversion
        pBuffer = new char[multiByteLength];

        if(pBuffer == NULL)
        {
            RpcPrintMessage(L"Out of Memory");
            hr = E_OUTOFMEMORY;
            goto end;
        }
    }

	// Perform the conversion to UTF8
    multiByteLength = WideCharToMultiByte(CP_UTF8, 0, strData, (int)data.length() + 1, pBuffer, multiByteLength, NULL, NULL);

    // Send Length
    CHK(SendDataLengthInBytes(multiByteLength));

    // Send Content
    CHK(SendDataContent(pBuffer, multiByteLength));

#ifdef _DEBUG
    size_t bufferLength = 1024 + data.length();
    wchar_t* buffer2 = new wchar_t[bufferLength];
    swprintf_s(buffer2, bufferLength, L"Wrote string: %s\n", data.data());
    RpcPrintMessage(buffer2);
    delete[] buffer2;
#endif

end:
    if(pBuffer != buffer)
    {
        delete [] pBuffer;
    }

    // TODO: Need to recover gracefully from sends that fail
    // but the connection is still available. i.e. Message too large or outgoing buffer is full.
    if(FAILED(hr))
    {
        RpcPrintMessage(L"WinsockStringSocket::SendData");
        RpcPrintError(hr);
        RaiseSocketDisconnectedEvent(m_clientDisconnectedCallback, this);
    }

    return hr;
}

// Turn compile optimization off for this function in order to use the Mock functionality in tests
// The reason we have created these RecvImpl / SendImpl functions is so we can mock recv in a different 
// Way when called to receive the length vs. receive the content
// Mock10 framework has limitations on filtered Mocking (Mock x only when called by y) when code is optimized so we had to take this approach
// Turn compile optimization off for this function in order to use the Mock functionality in tests
#pragma optimize( "", off)
int WinsockStringSocket::ReadStringRecvImpl(char* buffer, int strLength)
{
    return ::recv(m_socket, buffer, strLength, MSG_WAITALL);
}

bool WinsockStringSocket::ReadString(int strLength)
{
    static const size_t BUFFER_SIZE = 512;

    char buffer[BUFFER_SIZE];
    char * pBuffer = buffer;

    wchar_t wBuffer[BUFFER_SIZE];
    wchar_t *pWBuffer = wBuffer;

    bool success = true;

    if(strLength < 0)
    {
        wchar_t debugMsgBuffer[128];
        swprintf_s(debugMsgBuffer, 128, L"Invalid string length: %d", strLength);
        RpcPrintMessage(debugMsgBuffer);
        success = false;
        goto end;
    }
    else if((strLength + 1) > BUFFER_SIZE) // String length plus possible null-terminated character
    {
        pBuffer = new char[strLength + 1];
    }

    int result = ReadStringRecvImpl(pBuffer, strLength);
    if (result == 0)
    {
        RpcPrintMessage(L"socket was closed");
        success = false;
        goto end;
    }
    else if (result == SOCKET_ERROR)
    {
        RpcPrintMessage(L"recv failed with error");
        RpcPrintError(HRESULT_FROM_WSA);
        success = false;
        goto end;
    }

    if(result != strLength)
    {
        RpcPrintMessage(L"string length not met");
        success = false;
        goto end;
    }

    // String received may not be null-terminated so we append it just in case to 
    // make calls to MultiByteToWideChar more clear.
    pBuffer[strLength] = '\0';

    int len = MultiByteToWideChar(CP_UTF8, 0, pBuffer, -1, 0, 0); 

    if(len > BUFFER_SIZE)
    {
        pWBuffer = new wchar_t[len];
    }

    len = MultiByteToWideChar(CP_UTF8, 0, pBuffer, -1, pWBuffer, len);

    if(m_dataReceivedCallback != nullptr)
    {
		std::wstring incomingString(pWBuffer, len - 1);

        std::function<void()> taskFunction = [=]()
        {
            m_dataReceivedCallback(this, incomingString);
        };

        if (m_dataReceivedTaskInitialized)
        {
            // Chain the callbacks so they will be called in a chronological order for correct data flow
            m_dataReceivedCallbackTask = m_dataReceivedCallbackTask.then(taskFunction);
        }
        else
        {
            m_dataReceivedCallbackTask = create_task(taskFunction);
            m_dataReceivedTaskInitialized = true;
        }
    }

end:
    if(pWBuffer != wBuffer)
    {
        delete [] pWBuffer;
    }

    if(pBuffer != buffer)
    {
        delete [] pBuffer;
    }

    return success;
}

void WinsockStringSocket::RaiseSocketDisconnectedEvent(SocketDisconnectedCallback callback, WinsockStringSocket * socket)
{
    if(callback != nullptr)
    {
        create_task([=]()
        {
            callback(socket);
        });
    }
}

void WinsockStringSocket::WaitForDataReceivedCallbackTasks()
{
    if (m_dataReceivedTaskInitialized)
    {
        try
        {
            m_dataReceivedCallbackTask.wait();
        }
        catch(concurrency::invalid_operation& e)
        {
            RpcPrintMessage(e.what());
        }
    }
}

DWORD WinsockStringSocket::ThreadStart(LPVOID instance)
{
    return ((WinsockStringSocket*)instance)->Process();
}

void WinsockStringSocket::CloseSocket()
{
    if(m_socket != INVALID_SOCKET)
    {
        RpcPrintMessage(L"Closing socket");

        auto closeEvent = WSACreateEvent();

        // Attempt to gracefully close the connection
        // if for some reason this fails we still want to call
        // closesocket.
        if(closeEvent != nullptr)
        {
            WSAEventSelect(m_socket, closeEvent, FD_CLOSE);
            shutdown(m_socket, SD_SEND);
            WaitForSingleObject(closeEvent, 5000);			
            WSACloseEvent(closeEvent);
        }

        closesocket(m_socket);
        m_socket = INVALID_SOCKET;
    }
}

#endif // _WINRT_DLL